"""
参数优化工具模块
提供策略参数的网格搜索优化功能
"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from itertools import product
from tqdm import tqdm
import time
from datetime import datetime

# 设置中文字体支持
plt.rcParams['font.sans-serif'] = ['Arial Unicode MS', 'SimHei']
plt.rcParams['axes.unicode_minus'] = False
plt.style.use('seaborn-v0_8')
sns.set_palette("husl")

class ParameterOptimizer:
    def __init__(self, strategy_class, initial_capital=100000):
        """
        初始化参数优化器
        
        参数:
            strategy_class: 策略类（如MultiFactorMACDStrategy）
            initial_capital: 初始资金
        """
        self.strategy_class = strategy_class
        self.initial_capital = initial_capital
        self.optimization_results = []
        self.best_params = None
        self.best_performance = -np.inf
        
    def create_parameter_grid(self, complexity='medium'):
        """
        创建不同复杂度的参数网格
        
        参数:
            complexity: 参数网格复杂度 ('simple', 'medium', 'advanced')
        
        返回:
            param_grid: 参数网格字典
        """
        if complexity == 'simple':
            # 简单参数网格 - 基础MACD和RSI参数
            param_grid = {
                'fast_period': [12],
                'slow_period': [26],
                'signal_period': [9],
                'rsi_period': [14],
                'rsi_low_threshold': [30],
                'rsi_high_threshold': [70],
                'ma_short_period': [5],
                'ma_long_period': [20],
                'stop_loss_multiplier': [2.0],
                'take_profit_multiplier': [3.0]
            }
        elif complexity == 'medium':
            # 中等复杂度参数网格 - 适度调整关键参数
            param_grid = {
                'fast_period': [10, 12, 14],
                'slow_period': [20, 26, 30],
                'signal_period': [8, 9, 10],
                'rsi_period': [12, 14, 16],
                'rsi_low_threshold': [25, 30, 35],
                'rsi_high_threshold': [65, 70, 75],
                'ma_short_period': [5, 10],
                'ma_long_period': [20, 30],
                'stop_loss_multiplier': [1.5, 2.0, 2.5],
                'take_profit_multiplier': [2.5, 3.0, 3.5]
            }
        elif complexity == 'advanced':
            # 高级参数网格 - 广泛调整所有参数
            param_grid = {
                'fast_period': [8, 10, 12, 14, 16],
                'slow_period': [18, 20, 26, 30, 34],
                'signal_period': [7, 8, 9, 10, 11],
                'rsi_period': [9, 12, 14, 18, 21],
                'rsi_low_threshold': [20, 25, 30, 35, 40],
                'rsi_high_threshold': [60, 65, 70, 75, 80],
                'ma_short_period': [3, 5, 10],
                'ma_long_period': [20, 30, 50],
                'stop_loss_multiplier': [1.0, 1.5, 2.0, 2.5, 3.0],
                'take_profit_multiplier': [2.0, 2.5, 3.0, 3.5, 4.0]
            }
        else:
            raise ValueError("复杂度必须是 'simple', 'medium' 或 'advanced'")
        
        return param_grid
    
    def calculate_combination_count(self, param_grid):
        """
        计算参数组合数量
        
        参数:
            param_grid: 参数网格字典
        
        返回:
            count: 参数组合数量
        """
        count = 1
        for values in param_grid.values():
            count *= len(values)
        return count
    
    def run_grid_search(self, df, param_grid, performance_metric='sharpe_ratio', 
                       use_progress_bar=True, save_results=False, results_filename=None):
        """
        执行网格搜索参数优化
        
        参数:
            df: 股票数据DataFrame
            param_grid: 参数网格字典
            performance_metric: 性能指标 ('sharpe_ratio', 'total_return', 'calmar_ratio')
            use_progress_bar: 是否显示进度条
            save_results: 是否保存优化结果
            results_filename: 结果文件名
        
        返回:
            best_params: 最佳参数组合
            best_performance: 最佳性能值
            all_results: 所有参数组合的性能结果
        """
        # 记录开始时间
        start_time = time.time()
        
        # 清空之前的结果
        self.optimization_results = []
        self.best_params = None
        self.best_performance = -np.inf
        
        # 生成参数组合
        param_names = list(param_grid.keys())
        param_values = list(param_grid.values())
        total_combinations = self.calculate_combination_count(param_grid)
        
        print(f"开始参数优化，共 {total_combinations} 个参数组合...")
        print(f"性能评估指标: {performance_metric}")
        
        # 遍历所有参数组合
        combinations = product(*param_values)
        if use_progress_bar:
            combinations = tqdm(combinations, total=total_combinations, desc="优化进度")
        
        for params in combinations:
            param_dict = dict(zip(param_names, params))
            
            try:
                # 创建策略实例
                strategy = self.strategy_class(**param_dict)
                
                # 计算技术指标
                df_with_indicators = df.copy()
                df_with_indicators = strategy.calculate_macd(df_with_indicators)
                df_with_indicators = strategy.calculate_rsi(df_with_indicators)
                df_with_indicators = strategy.calculate_ma_cross(df_with_indicators)
                df_with_indicators = strategy.calculate_atr(df_with_indicators)
                df_with_indicators = strategy.calculate_volume_filter(df_with_indicators)
                
                # 生成交易信号
                df_with_signals = strategy.generate_signals(df_with_indicators)
                
                # 回测
                backtest_results = strategy.backtest_strategy(df_with_signals, self.initial_capital)
                
                # 获取性能指标
                if performance_metric == 'sharpe_ratio':
                    performance = backtest_results['sharpe_ratio']
                elif performance_metric == 'total_return':
                    performance = backtest_results['total_return']
                elif performance_metric == 'calmar_ratio':
                    performance = backtest_results['calmar_ratio']
                else:
                    performance = backtest_results['sharpe_ratio']
                
                # 保存结果
                result = {
                    'params': param_dict,
                    'performance': performance,
                    'sharpe_ratio': backtest_results['sharpe_ratio'],
                    'total_return': backtest_results['total_return'],
                    'annual_return': backtest_results['annual_return'],
                    'max_drawdown': backtest_results['max_drawdown'],
                    'calmar_ratio': backtest_results['calmar_ratio'],
                    'trade_count': backtest_results['trade_count'],
                    'win_rate': backtest_results['win_rate'],
                    'profit_loss_ratio': backtest_results['profit_loss_ratio']
                }
                
                self.optimization_results.append(result)
                
                # 更新最佳结果
                if performance > self.best_performance:
                    self.best_performance = performance
                    self.best_params = param_dict
                    
                    # 在进度条上显示当前最佳结果
                    if use_progress_bar:
                        combinations.set_postfix({
                            'best_score': f"{self.best_performance:.2f}",
                            'metric': performance_metric
                        })
            except Exception as e:
                print(f"参数组合 {param_dict} 出错: {e}")
                continue
        
        # 按性能指标排序
        self.optimization_results.sort(key=lambda x: x['performance'], reverse=True)
        
        # 计算优化耗时
        end_time = time.time()
        elapsed_time = end_time - start_time
        
        print(f"参数优化完成，耗时: {elapsed_time:.2f} 秒")
        print(f"最佳参数组合: {self.best_params}")
        print(f"最佳{performance_metric}: {self.best_performance:.4f}")
        
        # 保存结果
        if save_results:
            self.save_optimization_results(results_filename)
        
        return self.best_params, self.best_performance, self.optimization_results
    
    def get_best_params(self):
        """
        获取最佳参数组合
        
        返回:
            best_params: 最佳参数组合字典
        """
        if self.best_params is None:
            raise ValueError("请先运行参数优化")
        return self.best_params
    
    def get_optimization_results(self):
        """
        获取所有优化结果
        
        返回:
            optimization_results: 所有参数组合的优化结果列表
        """
        if not self.optimization_results:
            raise ValueError("请先运行参数优化")
        return self.optimization_results
    
    def save_optimization_results(self, filename=None):
        """
        保存参数优化结果到CSV文件
        
        参数:
            filename: 结果文件名，默认为当前时间戳
        """
        if not self.optimization_results:
            raise ValueError("没有优化结果可保存")
        
        # 生成默认文件名
        if filename is None:
            timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
            filename = f'optimization_results_{timestamp}.csv'
        
        # 转换结果为DataFrame
        results_df = pd.DataFrame()
        
        # 提取参数列
        param_cols = list(self.optimization_results[0]['params'].keys())
        for col in param_cols:
            results_df[col] = [r['params'][col] for r in self.optimization_results]
        
        # 添加性能指标列
        results_df['performance'] = [r['performance'] for r in self.optimization_results]
        results_df['sharpe_ratio'] = [r['sharpe_ratio'] for r in self.optimization_results]
        results_df['total_return'] = [r['total_return'] for r in self.optimization_results]
        results_df['annual_return'] = [r['annual_return'] for r in self.optimization_results]
        results_df['max_drawdown'] = [r['max_drawdown'] for r in self.optimization_results]
        results_df['calmar_ratio'] = [r['calmar_ratio'] for r in self.optimization_results]
        results_df['trade_count'] = [r['trade_count'] for r in self.optimization_results]
        results_df['win_rate'] = [r['win_rate'] for r in self.optimization_results]
        results_df['profit_loss_ratio'] = [r['profit_loss_ratio'] for r in self.optimization_results]
        
        # 保存到CSV
        results_df.to_csv(filename, index=False, encoding='utf-8-sig')
        print(f"参数优化结果已保存到: {filename}")
    
    def visualize_optimization_results(self, top_n=10, save_path=None):
        """
        可视化参数优化结果
        
        参数:
            top_n: 显示前N个最佳参数组合
            save_path: 图表保存路径
        """
        if not self.optimization_results:
            raise ValueError("没有优化结果可可视化")
        
        # 获取前N个最佳结果
        top_results = self.optimization_results[:top_n]
        
        # 创建图表
        fig, axes = plt.subplots(2, 2, figsize=(18, 15))
        
        # 1. 各参数组合的性能对比
        x_labels = [', '.join([f"{k}={v}" for k, v in r['params'].items()]) for r in top_results]
        x_labels = [label[:50] + '...' if len(label) > 50 else label for label in x_labels]
        
        performance_values = [r['performance'] for r in top_results]
        
        axes[0, 0].barh(x_labels, performance_values, color='skyblue')
        axes[0, 0].set_xlabel('性能值')
        axes[0, 0].set_ylabel('参数组合')
        axes[0, 0].set_title(f'前{top_n}个最佳参数组合性能对比')
        axes[0, 0].grid(True, alpha=0.3)
        
        # 2. 各性能指标散点图
        sharpe_ratios = [r['sharpe_ratio'] for r in self.optimization_results]
        total_returns = [r['total_return'] for r in self.optimization_results]
        max_drawdowns = [r['max_drawdown'] for r in self.optimization_results]
        
        scatter = axes[0, 1].scatter(total_returns, sharpe_ratios, 
                                   c=abs(np.array(max_drawdowns)), cmap='viridis', 
                                   alpha=0.7, s=50)
        axes[0, 1].set_xlabel('总收益率 (%)')
        axes[0, 1].set_ylabel('夏普比率')
        axes[0, 1].set_title('总收益率 vs 夏普比率 (颜色表示最大回撤绝对值)')
        axes[0, 1].grid(True, alpha=0.3)
        cbar = plt.colorbar(scatter, ax=axes[0, 1])
        cbar.set_label('最大回撤绝对值 (%)')
        
        # 3. 参数分布热图
        # 选择最重要的几个参数进行可视化
        important_params = ['fast_period', 'slow_period', 'signal_period', 'rsi_period']
        important_params = [p for p in important_params if p in self.optimization_results[0]['params']]
        
        if len(important_params) >= 2:
            # 创建参数组合的性能热图
            param1 = important_params[0]
            param2 = important_params[1]
            
            # 构建热图数据
            param1_values = sorted(set([r['params'][param1] for r in self.optimization_results]))
            param2_values = sorted(set([r['params'][param2] for r in self.optimization_results]))
            
            heatmap_data = np.zeros((len(param1_values), len(param2_values)))
            heatmap_data.fill(np.nan)  # 初始化为NaN
            
            for r in self.optimization_results:
                p1_idx = param1_values.index(r['params'][param1])
                p2_idx = param2_values.index(r['params'][param2])
                # 如果有多个相同参数组合，取性能最好的
                if np.isnan(heatmap_data[p1_idx, p2_idx]) or r['performance'] > heatmap_data[p1_idx, p2_idx]:
                    heatmap_data[p1_idx, p2_idx] = r['performance']
            
            # 绘制热图
            sns.heatmap(heatmap_data, annot=True, fmt='.2f', 
                        xticklabels=param2_values, yticklabels=param1_values,
                        cmap='RdYlGn', ax=axes[1, 0])
            axes[1, 0].set_xlabel(param2)
            axes[1, 0].set_ylabel(param1)
            axes[1, 0].set_title(f'{param1} vs {param2} 性能热图')
        else:
            axes[1, 0].text(0.5, 0.5, '参数不足，无法绘制热图', 
                          ha='center', va='center', transform=axes[1, 0].transAxes)
            axes[1, 0].axis('off')
        
        # 4. 交易统计指标分布
        trade_counts = [r['trade_count'] for r in self.optimization_results]
        win_rates = [r['win_rate'] for r in self.optimization_results]
        profit_loss_ratios = [r['profit_loss_ratio'] for r in self.optimization_results]
        
        # 过滤异常值
        profit_loss_ratios = [p if p < 5 else 5 for p in profit_loss_ratios]  # 限制盈亏比最大值为5
        
        axes[1, 1].hist2d(win_rates, profit_loss_ratios, bins=20, cmap='Blues')
        axes[1, 1].set_xlabel('胜率 (%)')
        axes[1, 1].set_ylabel('盈亏比')
        axes[1, 1].set_title('胜率 vs 盈亏比 分布')
        axes[1, 1].grid(True, alpha=0.3)
        
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            print(f"优化结果图表已保存到: {save_path}")
        
        plt.show()
    
    def analyze_parameter_sensitivity(self, param_name, save_path=None):
        """
        分析单个参数的敏感性
        
        参数:
            param_name: 参数名称
            save_path: 图表保存路径
        """
        if not self.optimization_results:
            raise ValueError("没有优化结果可分析")
        
        # 检查参数是否存在
        if param_name not in self.optimization_results[0]['params']:
            raise ValueError(f"参数 '{param_name}' 不存在于优化结果中")
        
        # 收集数据
        param_values = []
        performance_values = []
        sharpe_values = []
        max_drawdown_values = []
        trade_counts = []
        
        for r in self.optimization_results:
            param_values.append(r['params'][param_name])
            performance_values.append(r['performance'])
            sharpe_values.append(r['sharpe_ratio'])
            max_drawdown_values.append(r['max_drawdown'])
            trade_counts.append(r['trade_count'])
        
        # 创建DataFrame并按参数值排序
        sensitivity_df = pd.DataFrame({
            param_name: param_values,
            'performance': performance_values,
            'sharpe_ratio': sharpe_values,
            'max_drawdown': max_drawdown_values,
            'trade_count': trade_counts
        })
        
        # 按参数值分组，计算平均值
        sensitivity_df = sensitivity_df.groupby(param_name).mean().reset_index()
        sensitivity_df = sensitivity_df.sort_values(by=param_name)
        
        # 创建图表
        fig, axes = plt.subplots(3, 1, figsize=(12, 15))
        
        # 1. 参数值与性能关系
        axes[0].plot(sensitivity_df[param_name], sensitivity_df['performance'], 
                    marker='o', linewidth=2, markersize=8)
        axes[0].set_xlabel(param_name)
        axes[0].set_ylabel('性能值')
        axes[0].set_title(f'{param_name} 对性能的影响')
        axes[0].grid(True, alpha=0.3)
        
        # 2. 参数值与夏普比率、最大回撤关系
        axes[1].plot(sensitivity_df[param_name], sensitivity_df['sharpe_ratio'], 
                    marker='o', linewidth=2, markersize=8, label='夏普比率')
        axes[1].plot(sensitivity_df[param_name], sensitivity_df['max_drawdown'], 
                    marker='s', linewidth=2, markersize=8, label='最大回撤')
        axes[1].set_xlabel(param_name)
        axes[1].set_ylabel('指标值')
        axes[1].set_title(f'{param_name} 对风险调整收益的影响')
        axes[1].legend()
        axes[1].grid(True, alpha=0.3)
        
        # 3. 参数值与交易次数关系
        axes[2].plot(sensitivity_df[param_name], sensitivity_df['trade_count'], 
                    marker='^', linewidth=2, markersize=8, color='green')
        axes[2].set_xlabel(param_name)
        axes[2].set_ylabel('交易次数')
        axes[2].set_title(f'{param_name} 对交易频率的影响')
        axes[2].grid(True, alpha=0.3)
        
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            print(f"参数敏感性分析图表已保存到: {save_path}")
        
        plt.show()

# 使用示例
def example_optimization():
    """
    参数优化使用示例
    """
    from macd_multi_factor_strategy import MultiFactorMACDStrategy
    
    # 创建模拟数据
    print("创建模拟数据...")
    dates = pd.date_range('2023-01-01', '2023-12-31', freq='D')
    np.random.seed(42)
    
    # 生成模拟股价数据（带趋势）
    trend = np.linspace(100, 130, len(dates))
    noise = np.random.randn(len(dates)) * 2
    price = trend + np.cumsum(noise) * 0.4
    price = np.maximum(price, 1)  # 确保价格为正
    
    df = pd.DataFrame({
        '日期': dates,
        '股票代码': '000001.SZ',
        '开盘': price * (1 - np.random.uniform(0, 0.02, len(dates))),
        '收盘': price,
        '最高': price * (1 + np.random.uniform(0, 0.03, len(dates))),
        '最低': price * (1 - np.random.uniform(0, 0.03, len(dates))),
        '成交量': np.random.randint(1000000, 8000000, len(dates)),
        '成交额': price * np.random.randint(1000000, 8000000, len(dates)),
        '振幅': np.random.uniform(1, 6, len(dates)),
        '涨跌幅': np.random.uniform(-4, 4, len(dates)),
        '涨跌额': np.random.uniform(-3, 3, len(dates)),
        '换手率': np.random.uniform(0.8, 4, len(dates))
    })
    
    # 创建优化器实例
    print("初始化参数优化器...")
    optimizer = ParameterOptimizer(MultiFactorMACDStrategy, initial_capital=100000)
    
    # 创建简单的参数网格（快速测试）
    param_grid = optimizer.create_parameter_grid(complexity='simple')
    print(f"参数网格: {param_grid}")
    print(f"参数组合数量: {optimizer.calculate_combination_count(param_grid)}")
    
    # 运行网格搜索
    print("运行参数优化...")
    best_params, best_performance, all_results = optimizer.run_grid_search(
        df, param_grid, performance_metric='sharpe_ratio', 
        use_progress_bar=True, save_results=True
    )
    
    # 可视化优化结果
    print("可视化优化结果...")
    optimizer.visualize_optimization_results(top_n=5, save_path='optimization_results.png')
    
    # 分析参数敏感性
    if 'fast_period' in best_params:
        optimizer.analyze_parameter_sensitivity('fast_period', save_path='sensitivity_fast_period.png')
    
    return best_params, best_performance

if __name__ == "__main__":
    example_optimization()